from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import importlib
import os
import glob
import argparse
from dotenv import load_dotenv
from py_files.openai_proxy import router as openai_proxy_router
from py_files.prompt_agent_generate.image_agent_generate import image_agent_router
# 加载环境变量
load_dotenv()

# 创建FastAPI应用
app = FastAPI(
    title="LiteLLM API",
    description="OpenAI兼容的API接口，支持加载多种大模型进行对话",
    version="1.0.0",
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 根路由
@app.get("/")
async def root():
    return {"message": "see api in /docs"}

app.include_router(openai_proxy_router)
app.include_router(image_agent_router)
# # 动态加载py_files目录下的所有路由
# py_files_path = os.path.join(os.path.dirname(__file__), "py_files")
# os.makedirs(py_files_path, exist_ok=True)

# # 查找所有路由文件并导入
# router_files = glob.glob(os.path.join(py_files_path, "*.py"))
# for router_file in router_files:
#     file_name = os.path.basename(router_file)
#     module_name = f"py_files.{file_name[:-3]}"  # 移除.py后缀
#     try:
#         module = importlib.import_module(module_name)
#         if hasattr(module, 'router'):
#             app.include_router(module.router)
#     except ImportError as e:
#         print(f"Error importing {module_name}: {e}")

# 主函数
if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="LiteLLM API服务")
    parser.add_argument("--test-models", action="store_true", help="启动前测试所有模型连通性")
    args = parser.parse_args()
    
    # 如果指定了--test-models参数，则测试模型连通性
    if args.test_models:
        print("测试模型连通性...")
        try:
            from py_files.config import test_all_models
            available, unavailable, _ = test_all_models(verbose=True)
            print(f"可用模型: {', '.join(available)}")
            if unavailable:
                print(f"不可用模型: {', '.join(unavailable)}")
        except Exception as e:
            print(f"测试模型时发生错误: {e}")
    else:
        print("已跳过模型连通性测试。如需测试，请使用 --test-models 参数")
    
    # 启动服务
    uvicorn.run("main:app", host="0.0.0.0", port=8007, reload=True) 